import datetime
from app.db.models.tasks import Task
from app.db.models.wass import ScanEvents


class AssetChangeTaskAddition:

    def service(self):
        print(f"{datetime.datetime.now()} 开始同步数据......")

        task_session_ids = set()
        task_session_maps = dict()
        tasks = Task.objects.filter(status="completed", taskType="asset").only("taskSessionId", "result.addition")
        for task in tasks:
            task_session_id = task.taskSessionId
            if task_session_id:
                task_session_ids.add(task_session_id)
                task_session_maps[task_session_id] = {"addition": task.result.addition, "_id": task.id}

        print(f"{datetime.datetime.now()} 需要同步的任务数量: {len(task_session_ids)}")

        count = 0
        scan_events = ScanEvents.objects.filter(task_session_id__in=list(task_session_ids)).only("task_session_id",
                                                                                                 "detail")
        for scan_event in scan_events:
            task = task_session_maps.get(scan_event.task_session_id)
            if task:
                detail = scan_event.detail
                addition = task.get("addition", {})
                addition.update({
                    "waf_list": detail.get("waf_list", []),
                    "alexa": detail.get("alexa", ''),
                    "basic_info": detail.get("basic_info", {"server": "", "title": ""}),
                    "icp": detail.get("icp", {"desc": "未查到"}),
                    "ipdb": detail.get("ipdb", {"country_name": "", "region_name": "", "city_name": ""}),
                    "subdomain": detail.get("subdomain", {}),
                    "http_get": detail.get("http_get", {"status": 0.0}),
                    "http_method": detail.get("http_method", ""),
                    "port": detail.get("port", {}),
                })
                Task.objects(pk=task.get("_id")).update(result__addition=addition)
                count += 1

        print(f"完成同步的任务数量: {count}")
        print(f"{datetime.datetime.now()} 结束！！！")
